package main

import (
	"context"
	"io/ioutil"
	"github.com/chromedp/cdproto/page"
	"github.com/chromedp/chromedp"
	_ "errors"
	"fmt"
	"time"
	"os"
	"strings"
	"runtime"

	"flag"
	"github.com/remeh/sizedwaitgroup"
	_ "golang.org/x/text/runes"
)

const (
	host = "http://localhost:8080/"

	// Answer试卷宽，单位英寸
	answerPaperWidth = 3.7465
	// Answer上边距，单位英寸
	answerMarginTop = 0.3937
	// Answer上边距，单位英寸
	answerMarginBottom = 0.3937
	// Answer左边距，单位英寸
	answerMarginLeft = 0.3937
	// Answer右边距，单位英寸
	answerMarginRight = 0.3937

	//上边距，单位英寸
	marginTop = 0.5906
	//上边距，单位英寸
	marginBottom = 0.5906
	//左边距，单位英寸
	marginLeft = 0.5906
	//右边距，单位英寸
	marginRight = 0.5906
	//等待加载时长
	sleep = 3
)

func main() {
	var basePath,filesPath,chromePath string
	//html根目录
	basePath = "/Users/fangzhaocheng/tools/html"
	//扫描htm文件夹目录
	filesPath = "/Users/fangzhaocheng/tools/html/107159411469388801269/100159411358507200331/pdf/18班/数学"

	flag.StringVar(&basePath, "b", "/Users/fangzhaocheng/tools/html", "html根目录")
	flag.StringVar(&filesPath, "f", "", "扫描htm文件夹目录")
	flag.StringVar(&chromePath, "a", "", "chrome路径")
	flag.Parse()
	fmt.Printf("basePath：%s\nfilesPath：%s \n", basePath, filesPath)
	if len(basePath) == 0 {
		fmt.Errorf("html根目录不能为空")
		return
	}
	if len(filesPath) == 0 {
		fmt.Errorf("扫描htm文件夹目录")
		return
	}
	if len(chromePath) > 0 {
		chromedp.ExecPath(chromePath)
	}
	n := runtime.NumCPU()
	runtime.GOMAXPROCS(n - 1)
	// 获取文件列表
	files, err1 := GetAllFiles(filesPath)
	fmt.Println(len(files))
	swg := sizedwaitgroup.New(n*2)
	fmt.Println(time.Now())
	if err1 == nil {
		for _, fileStr := range files {
			swg.Add()
			go func(fileStr string) {
				defer swg.Done()
				generatePdf(fileStr,basePath)
			}(fileStr)
		}
	} else {
		fmt.Println(err1)
	}
	swg.Wait()
	fmt.Println(time.Now())
}

func generatePdf(fileStr ,basePath string) {
	htmlPath, pdfPath := getPdfPath(fileStr, basePath)
	htmlSrc := host + htmlPath
	err := ChromePrintPdf(htmlSrc, pdfPath)
	if err != nil {
		fmt.Println(err)
	}
	fmt.Println(pdfPath)
	os.Remove(fileStr)
}

func generatePdfOld() {
	host := "http://localhost:8080"
	basePath := "/Users/fangzhaocheng/tools/html"
	filesPath := "/Users/fangzhaocheng/tools/html/107159411469388801269/1111"
	files, err1 := GetAllFiles(filesPath)
	var htmlSrc string
	if err1 == nil {
		for _, fileStr := range files {
			htmlPath, pdfPath := getPdfPath(fileStr, basePath)
			htmlSrc = host + htmlPath
			err := ChromePrintPdf(htmlSrc, pdfPath)
			if err != nil {
				fmt.Println(err)
			}
			fmt.Println(pdfPath)

		}
	} else {
		fmt.Println(err1)
	}
}

func getPdfPath(filePath, basePath string) (htmlPath, pdfPathName string) {

	htmlPath = filePath[len(basePath):]
	comma := strings.Index(filePath, ".")
	pdfPathName = filePath[:comma] + ".pdf"
	return
}

func ChromePrintPdf(url string, to string) error {
	ctx, cancel := chromedp.NewContext(context.Background())
	defer cancel()

	var buf []byte
	err := chromedp.Run(ctx, chromedp.Tasks{
		chromedp.Navigate(url),
		//chromedp.WaitReady("body"),
		chromedp.Sleep(sleep * time.Second),
		chromedp.ActionFunc(func(ctx context.Context) error {
			var err error
			if strings.Index(to, "Answer") > 0 {
				buf, _, err = page.PrintToPDF().
					WithPaperWidth(answerPaperWidth).
					WithMarginTop(answerMarginTop).
					WithMarginBottom(answerMarginBottom).
					WithMarginLeft(answerMarginLeft).
					WithMarginRight(answerMarginRight).
					Do(ctx)
			} else {
				buf, _, err = page.PrintToPDF().
				//单位 英寸
					WithDisplayHeaderFooter(false).
					WithMarginTop(marginTop).
					WithMarginBottom(marginBottom).
					WithMarginLeft(marginLeft).
					WithMarginRight(marginRight).
					Do(ctx)
			}

			return err
		}),
	})
	if err != nil {
		return fmt.Errorf("chromedp Run failed,err:%+v", err)
	}

	if err := ioutil.WriteFile(to, buf, 0644); err != nil {
		return fmt.Errorf("write to file failed,err:%+v", err)
	}

	return nil
}

//获取指定目录下的所有文件,包含子目录下的文件
func GetAllFiles(dirPth string) (files []string, err error) {
	var dirs []string
	dir, err := ioutil.ReadDir(dirPth)
	if err != nil {
		return nil, err
	}

	PthSep := string(os.PathSeparator)

	for _, fi := range dir {
		if fi.IsDir() { // 目录, 递归遍历
			dirs = append(dirs, dirPth+PthSep+fi.Name())
			GetAllFiles(dirPth + PthSep + fi.Name())
		} else {
			// 过滤指定格式
			ok := strings.HasSuffix(fi.Name(), ".htm")
			if ok {
				files = append(files, dirPth+PthSep+fi.Name())
			}
		}
	}

	// 读取子目录下文件
	for _, table := range dirs {
		temp, _ := GetAllFiles(table)
		for _, temp1 := range temp {
			files = append(files, temp1)
		}
	}

	return files, nil
}
